import random
import torch as T
import torch.nn.functional as F
from MADDPG.agent import Agent
import numpy as np

T.autograd.set_detect_anomaly(True)


class MADDPG:
    def __init__(self, actor_dims, critic_dims, n_agents, n_actions, num_users, noise_type, max_episodes,
                 scenario='simple', alpha=0.0005, beta=0.001, chkpt_dir='tmp/maddpg/'):  # chkpt_dir是存储路径
        self.curr_step = 0
        self.agents = []
        self.num_user = num_users
        self.n_agents = n_agents
        self.n_actions = n_actions
        self.max_episodes = max_episodes
        chkpt_dir += scenario
        for agent_idx in range(self.n_agents):
            self.agents.append(Agent(actor_dims[agent_idx], critic_dims,
                                     n_actions, n_agents, num_users, agent_idx, noise_type, chkpt_dir=chkpt_dir,
                                     alpha=alpha, beta=beta, fc1=128, fc2=128, fc3=128, gamma=0.95, tau=0.01))

    def save_checkpoint(self):  # 保存模型状态，用于调试和回复
        print('... saving checkpoint ...')
        for agent in self.agents:
            agent.save_models()

    def load_checkpoint(self):  # 加载模型状态，用于恢复训练
        print('... loading checkpoint ...')
        for agent in self.agents:
            agent.load_models()

    def choose_action(self, obs, e):
        self.curr_step += 1
        actions = []
        for agent_idx, agent in enumerate(self.agents):
            action = agent.choose_action_noised(obs[agent_idx], e, self.max_episodes)
            actions.append(action)
        return actions

    # 数据处理-critic更新-actor更新-反向传播与优化
    def learn(self, memory):
        if not memory.ready():
            return
        actor_states, states, actions, rewards, actor_states_, states_, dones = memory.sample_buffer()

        device = self.agents[0].actor.device
        # 将不同的数据转换为PyTorch张量并移动到指定的设备
        states = T.tensor(states, dtype=T.float).to(device)
        actions = np.array(actions)  # 将列表转换为 numpy 数组
        actions = T.tensor(actions, dtype=T.float).to(device)
        rewards = T.tensor(rewards, dtype=T.float).to(device)
        states_ = T.tensor(states_, dtype=T.float).to(device)
        dones = T.tensor(dones).to(device)
        all_agents_new_actions = []
        old_agents_actions = []

        for agent_idx, agent in enumerate(self.agents):
            new_states = T.tensor(actor_states_[agent_idx],
                                  dtype=T.float).to(device)

            new_pi = agent.target_actor.forward(new_states)

            all_agents_new_actions.append(new_pi)
            old_agents_actions.append(actions[agent_idx])

        new_actions = T.cat([acts for acts in all_agents_new_actions], dim=1)
        # for agent_idx, acts in enumerate(all_agents_new_actions):
        #     print(f"Shape of actions for agent {agent_idx}: {acts.shape}")
        old_actions = T.cat([acts for acts in old_agents_actions], dim=1)
        # print(f"Shape of new_actions: {new_actions.shape}")
        for agent_idx, agent in enumerate(self.agents):
            #计算critic损失——输入是状态和所有动作，输出q值
            with T.no_grad():
                critic_value_ = agent.target_critic.forward(states_, new_actions).flatten()
                # 打印 critic_value_ 的形状和部分数据
                # print(f"Shape of critic_value_: {critic_value_.shape}")
                # print(f"Critic value from target_critic (critic_value_): {critic_value_[:5]}")  # 打印前5个值
                target = rewards[:, agent_idx] + (1 - dones[:, 0].int()) * agent.gamma * critic_value_

            critic_value = agent.critic.forward(states, old_actions).flatten()
            # # 打印 critic_value 的形状和部分数据
            # print(f"Shape of critic_value: {critic_value.shape}")
            # print(f"Critic value from current critic (critic_value): {critic_value[:5]}")  # 打印前5个值
            critic_loss = F.mse_loss(target, critic_value)
            agent.critic.optimizer.zero_grad()
            critic_loss.backward(retain_graph=True)
            agent.critic.optimizer.step()

            # 计算actor损失——输入是每个agent的观测，输出每个agent的动作
            mu_states = T.tensor(actor_states[agent_idx], dtype=T.float).to(device)  #当前agent的状态
            oa = old_actions.clone()
            oa[:, agent_idx * self.n_actions:agent_idx * self.n_actions + self.n_actions] = agent.actor.forward(
                mu_states)
            actor_loss = -T.mean(agent.critic.forward(states, oa).flatten())
            agent.actor.optimizer.zero_grad()
            actor_loss.backward(retain_graph=True)
            agent.actor.optimizer.step()
            agent.update_network_parameters()
